// "Copyright [2021] <Copyright Shengkai Lin>"

#include "gpu_schedule_algo.h"

using namespace omnetpp;

namespace ai_sim {
void SequentialScheduleAlgorithm::scheduleGpu(
    taskboard& tmp_taskboard, TopoManager * topos, int job_num) {
        int total_workers = topos->getWorkersTopo()->getNumNodes();
        int first_ptr = 0;
        while(true) {
            if (first_ptr >= total_workers) {
                break;
            }
            topoNode *tmp_node_ptr = topos->getWorkersTopo()->getNode(first_ptr);
            if (tmp_node_ptr->isEnabled()) {
                break;
            }
            first_ptr++;
        }
        int ptr_index = first_ptr;
        vector<int> valid_ptr;
        while(true) {
            if(ptr_index >= total_workers) {
                break;
            }
            topoNode *tmp_node_ptr = topos->getWorkersTopo()->getNode(ptr_index);
            if (tmp_node_ptr->isEnabled()) {
                valid_ptr.push_back(ptr_index);
            }
            ptr_index+=job_num;
        }
        // std::cout<<endl;
        for (int t = 0; t < total_workers/job_num; t++) {
            int i = valid_ptr[t];
            // std::cout<<i<<endl;
            vector<int> temp_list;
            topoNode *node_ptr = topos->getWorkersTopo()->getNode(i);
            // EV<< node_ptr->getModule()->getModuleType();
            if (node_ptr->isEnabled()) {
                tmp_taskboard.used_gpus.push_back(node_ptr);
                node_ptr->disable();
                temp_list.push_back(i);
            }
            if (tmp_taskboard.used_gpus.size() >= tmp_taskboard.the_task.numGpus) {
                EV_INFO << "Total gpus " << total_workers << std::endl;
                EV_INFO << "Need gpus " <<
                    tmp_taskboard.the_task.numGpus << std::endl;
                break;
            }
        }
        if (tmp_taskboard.used_gpus.size() < tmp_taskboard.the_task.numGpus) {
            std::cout << "used gpus" << tmp_taskboard.used_gpus.size()  << std::endl;
            std::cout << "Total gpus" << total_workers  << std::endl;
            std::cout << "Need gpus" << tmp_taskboard.the_task.numGpus << std::endl;
            throw cRuntimeError("DONT HAVE ENOUGH GPUS");
            return;
        }
}

}  // namespace ai_sim
